package ai.neuralnet.builders;

import ai.neuralnet.NNBuilder;
import ai.neuralnet.NeuralNetwork;
import ai.neuralnet.TNeuron;

import java.util.ArrayList;

public class OverlappingPyramidBuilder<N extends TNeuron, NN extends NeuralNetwork<N>> extends NNBuilder<N, NN>
{
    private int[] layerSizes = new int[]{50, 10, 2, 1};

    private int layersToOutput = 3;

    private double percentOverlap = 0.1;

    private int[] logicalGroups = new int[]{25, 1, 1, 1};


    public OverlappingPyramidBuilder(Class<N> nClass, Class<NN> nnClass)
    {
        super(nClass, nnClass);
    }

    @Override
    public NN build()
    {
        NN network = makeNeuralNetwork();

        ArrayList<N> inputNeurons = new ArrayList<N>(this.inputCount);

        for (int i = 0; i < this.inputCount; i++)
        {
            inputNeurons.add(makeNeuron(1));
            inputNeurons.get(i).setIsInputNeuron(true);
        }

        N[] inputArray = makeNeuronArray(this.inputCount);
        inputNeurons.toArray(inputArray);
        network.setInputNeurons(inputArray);

        ArrayList<N> prevLayer = inputNeurons;
        ArrayList<ArrayList<N>> hiddenLayers = new ArrayList<ArrayList<N>>();

        int prevSize = this.inputCount;
        for (int index = 0; index < layerSizes.length; index++)
        {
            ArrayList<N> hiddenLayer = new ArrayList<N>();
            int layerSize = layerSizes[index];
            // number of inputCount per neuron before overlap
            int div = (int) Math.max(Math.floor(prevSize / (double) (layerSize * logicalGroups[index])), 1) * logicalGroups[index];
            //layer size must be upper bounded
            layerSize = Math.min(layerSize, prevSize / div);

            //number of additional neurons to add on either side of the section from overlap
            int overlap = (int) Math.ceil(percentOverlap * div / (2 * logicalGroups[index])) * logicalGroups[index];
            for (int i = 0; i < layerSize; i++)
            {

                int start = Math.max(0, i * div - overlap);
                int end = Math.min(prevSize, (i + 1) * div + overlap);
                N n = makeNeuron(end - start);

                for (int j = start; j < end; j++)
                {
                    n.addInputNeuron(prevLayer.get(j));
                    prevLayer.get(j).addOutputNeuron(n);
                }
                hiddenLayer.add(n);
            }

            network.addHiddenNeuron(hiddenLayer);

            hiddenLayers.add(hiddenLayer);

            prevLayer = hiddenLayer;

            prevSize = layerSize;
        }

        N[] outNeurons = makeNeuronArray(outputCount);
        int hLayers = layerSizes.length;
        //count the number of inputs to the output layer
        int outNeurInputs = 0;
        for (int j = hLayers - 1; j >= hLayers - layersToOutput; j--)
        {
            outNeurInputs += layerSizes[j];
        }


        for (int i = 0; i < outputCount; i++)
        {
            N n = makeNeuron(outNeurInputs);

            for (int j = hLayers - 1; j >= hLayers - layersToOutput; j--)
            {
                n.addInputNeuron(hiddenLayers.get(j));
                for (N in : hiddenLayers.get(j))
                {
                    in.addOutputNeuron(n);
                }
            }

            outNeurons[i] = n;
        }

        network.setOutputNeurons(outNeurons);

        return network;
    }

    public void setLayerSizes(int[] layerSizes)
    {
        this.layerSizes = layerSizes;
    }

    public void setPercentOverlap(double percentOverlap)
    {
        this.percentOverlap = percentOverlap;
    }

    public void setLogicalGroups(int[] logicalGroups)
    {
        this.logicalGroups = logicalGroups;
    }

    public void setLayersToOutput(int layersToOutput)
    {
        this.layersToOutput = layersToOutput;
    }
}
